import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor, Cursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.classic_strategy.turtle_strategy import TurtleTrade


plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
register_matplotlib_converters()


class Block():
    def __init__(self, codes, start, end):
        if isinstance(codes, list):
            self.codes = codes
            self.names = None
        else:
            self.codes = codes['code'].to_list()
            self.names = codes[['code', 'name']].set_index(['code'])
        self.struct = QA.QA_fetch_index_day_adv(self.codes, start, end)
        self.data = self.struct.data
        self.data_close = self.data['close'].unstack() / self.data['close'].unstack().iloc[0, :]

    def plot(self, codes=None):
        if codes is None:
            ax = self.data_close.plot()
        else:
            data = self.data_close.loc[:, codes]
            d = self.names['name'].to_dict()
            data = data.rename(columns=self.names['name'].to_dict())
            ax = data.plot()
        cursor = Cursor(ax, useblit=True, color='r', lw=1)
        plt.show()

    def best_returns(self, bar_count=30, N=5):
        # data = self.data_close.rename(columns=self.names['name'].to_dict())
        data = self.data_close
        data = data.iloc[-bar_count:, :]
        data = data / data.iloc[0, :]
        data.to_csv('data/block_return.csv', encoding='utf_8_sig')
        # print(data.head())
        # print(data.describe())
        des = data.describe().T
        des['best'] = des['mean'] / des['std']
        des = des.sort_values(by='max', ascending=False)
        # des.to_csv('data/block_return_sort.csv', encoding='utf_8_sig')
        des = des.iloc[:N, :]
        print(des)
        return des.index.to_list()


if __name__ == '__main__':
    code_list = QA.QA_fetch_index_list()
    # list = list.code.to_list()
    # 880301到880685  8803到8806开头都是行业
    blocks = code_list[code_list.code.str.startswith('000001')
                 | code_list.code.str.startswith('8803')
                 | code_list.code.str.startswith('8804')
                 | code_list.code.str.startswith('8805')
                 # | code_list.code.str.startswith('8806')
                 | code_list.code.str.startswith('000050')
                 | code_list.code.str.startswith('000300')]
    etf_list = QA.QA_fetch_etf_list()
    etfs = etf_list[etf_list.code.str.startswith('510')
                 | etf_list.code.str.startswith('512')
                 | etf_list.code.str.startswith('515')]
    # etfs = ['159919', '510300']
    # blocks = ['000001', '510050', '510300']
    blocks = ['512760', '512660', '512690', '515790', '512480', '515030']  # 512760 在9月10日拆分1:2
    start = '2020-01-01'
    end = '2021-01-26'
    # main(codes, start, end)
    blk = Block(blocks, start, end)
    blk.plot()
    # rtn = blk.best_returns(bar_count=30)
    # blk.plot(rtn)
